package nachos.threads;

import nachos.machine.*;

import java.util.TreeSet;
import java.util.HashSet;
import java.util.Iterator;
import java.util.List;
import java.util.ArrayList;
import java.util.Set;
import java.util.HashSet;

// TODO: 
//  * Test decreasePriority.

/**
 * A scheduler that chooses threads based on their priorities.
 *
 * <p>
 * A priority scheduler associates a priority with each thread. The next thread
 * to be dequeued is always a thread with priority no less than any other
 * waiting thread's priority. Like a round-robin scheduler, the thread that is
 * dequeued is, among all the threads of the same (highest) priority, the
 * thread that has been waiting longest.
 *
 * <p>
 * Essentially, a priority scheduler gives access in a round-robin fassion to
 * all the highest-priority threads, and ignores all other threads. This has
 * the potential to
 * starve a thread if there's always a thread waiting with higher priority.
 *
 * <p>
 * A priority scheduler must partially solve the priority inversion problem; in
 * particular, priority must be donated through locks, and through joins.
 */
public class PriorityScheduler extends Scheduler {

    /**
     * Allocate a new priority scheduler.
     */
    public PriorityScheduler() {
    }
    
    /**
     * Allocate a new priority thread queue.
     *
     * @param	transferPriority	<tt>true</tt> if this queue should
     *					transfer priority from waiting threads
     *					to the owning thread.
     * @return	a new priority thread queue.
     */
    public ThreadQueue newThreadQueue(boolean transferPriority) {
	return new PriorityQueue(transferPriority);
    }

    public int getPriority(KThread thread) {
	Lib.assertTrue(Machine.interrupt().disabled());
		       
	return getThreadState(thread).getPriority();
    }

    public int getEffectivePriority(KThread thread) {
	Lib.assertTrue(Machine.interrupt().disabled());
	return getThreadState(thread).getEffectivePriority();
    }

    public void setPriority(KThread thread, int priority) {
	Lib.assertTrue(Machine.interrupt().disabled());
		       
	Lib.assertTrue(priority >= priorityMinimum &&
		       priority <= priorityMaximum);

	getThreadState(thread).setPriority(priority);
    }

    /**
     * Increase the priority of the current thread, if
     * possible. Return boolean indicating whether or not we actually
     * increased the priority.
     */
    public boolean increasePriority() {
	boolean intStatus = Machine.interrupt().disable();
	KThread thread = KThread.currentThread();
	int priority = getPriority(thread);
	boolean canIncrease = priority < priorityMaximum;

	if (canIncrease) {
	    setPriority(thread, priority+1);
	}

	Machine.interrupt().restore(intStatus);
	return canIncrease;
    }

    /**
     * Decrease the priority of the current thread, if
     * possible. Return boolean indicating whether or not we actually
     * decreased the priority.
     */
    public boolean decreasePriority() {
	boolean intStatus = Machine.interrupt().disable();
	KThread thread = KThread.currentThread();
	int priority = getPriority(thread);
	boolean canDecrease = priority > priorityMinimum;
	
	if (canDecrease)
	    setPriority(thread, priority-1);

	Machine.interrupt().restore(intStatus);
	return canDecrease;
    }

    /**
     * The default priority for a new thread. Do not change this value.
     */
    public static final int priorityDefault = 1;
    /**
     * The minimum priority that a thread can have. Do not change this value.
     */
    public static final int priorityMinimum = 0;
    /**
     * The maximum priority that a thread can have. Do not change this value.
     */
    public static final int priorityMaximum = 7;    

    /**
     * Return the scheduling state of the specified thread.
     *
     * @param	thread	the thread whose scheduling state to return.
     * @return	the scheduling state of the specified thread.
     */
    protected ThreadState getThreadState(KThread thread) {
	if (thread.schedulingState == null)
	    thread.schedulingState = new ThreadState(thread);

	return (ThreadState) thread.schedulingState;
    }

    /**
     * A <tt>ThreadQueue</tt> that sorts threads by priority.
     */
    protected class PriorityQueue extends ThreadQueue {

	PriorityQueue(boolean transferPriority) {
	    this.transferPriority = transferPriority;
	}

	/**
	 * Return true if the thread at index i should be released
	 * before the thread at index j, false otherwise.
	 */
	public boolean greaterThan(int i, int j) {
	    return threads.get(i).compareTo(threads.get(j)) > 0;
	}

	/**
	 * Swap the threads at index i and j, updating the index
	 * property of the two threads.
	 */
	private void swap(int i, int j) {
	    ThreadState ti = threads.get(i);
	    ThreadState tj = threads.get(j);
	    ti.index = j;
	    tj.index = i;
	    threads.set(i, tj);
	    threads.set(j, ti);
	}

	/**
	 * Called by ThreadState to notify the queue that its priority
	 * has been increased. When a thread's priority increases, we
	 * simply need to bubble it up to its new position in the
	 * queue.
	 */
	void priorityIncreased(ThreadState ts) {
	    bubbleUp(ts);
	}

	/**
	 * Called by ThreadState to notify the queue that its priority
	 * has been decreased. Since max-heaps don't directly support
	 * a decrease-key operation, we handle this by temporarily
	 * raising the thread's priority above the maximum value so
	 * that it rises to the top of the queue, then popping it,
	 * restoring the priority to what it should be and adding it
	 * back onto the queue.
	 */
	void priorityDecreased(ThreadState t) {
	    // Temporarily increase t's priority to higher than the
	    // maximum so that it bubbles to the top of the heap.
	    int priority = t.getPriority();
	    t.priority = priorityMaximum + 1;
	    t.effectivePriority = priorityMaximum + 1;
	    bubbleUp(t);

	    // t should now be on top of the heap, so pop and make
	    // sure that we got t back.
	    Lib.assertTrue(pop() == t);

	    // Now add t back onto the heap by putting it at the
	    // bottom and then bubbling it up to its right spot.
	    t.priority = priority;
	    t.effectivePriority = priority;
	    threads.add(t);
	    t.index = last();
	    bubbleUp(last());

	    if (owner != null) {
		owner.calculateEffectivePriority();
	    }
	}

	void bubbleUp(ThreadState t) {
	    bubbleUp(t.index);
	}

	/** 
	 * Move the thread at index i up in the heap until it is no
	 * longer greater than its parent. If this process results in
	 * the specified thread rising to the top of the heap, we may
	 * need to donate its priority to the owner of the queue. This
	 * could cause that thread's priority to change, so
	 * recursively bubble it up to the top of whatever heap it is
	 * on.
	 */
	void bubbleUp(int i) {
	    while (i > 0 && greaterThan(i, parent(i))) {
		swap(i, parent(i));
		i = parent(i);
	    }

	    if (i == 0 &&
		owner != null &&
		transferPriority &&
		highestPriority() > owner.getEffectivePriority()) {
		owner.effectivePriority = highestPriority();
		if (owner.waitQueue != null) {
		    owner.waitQueue.bubbleUp(owner);
		}
	    }
	}

	/**
	 * Add the given thread to the end of the queue.
	 */
	public void waitForAccess(KThread thread) {
	    Lib.assertTrue(Machine.interrupt().disabled());
	    
	    ThreadState ts = getThreadState(thread);

	    // If this thread currently owns this queue, then release
	    // the ownership since this thread is moving to the back
	    // of the line.
	    if (owner == ts)
		owner = null;

	    // Add the thread to the end of the queue, set its
	    // waitQueue, index, and entry time, and then move it up
	    // to its proper position in the heap.
	    threads.add(ts);
	    ts.waitForAccess(this, last(), time++);
	    bubbleUp(last());
	}

	/**
	 * Sets my owner to the given thread and calls owner(this) on
	 * that thread, in order to add this queue to its list of
	 * owned queues and clear out its wait queue.
	 */
	public void acquire(KThread thread) {
	    Lib.assertTrue(Machine.interrupt().disabled());
	    owner = getThreadState(thread);
	    owner.acquire(this);
	}

	/**
	 * Pop the next thread of the queue and return it.
	 */
	private ThreadState pop() {
	    swap(0, last());
	    ThreadState result = threads.remove(last());
	    heapify(0);
	    return result;
	}

	
	public KThread nextThread() {

	    Lib.assertTrue(Machine.interrupt().disabled());

	    // If a thread owns this queue, break that ownership.
	    if (owner != null) {
		owner.ownedQueues.remove(this);
		owner.calculateEffectivePriority();
		owner = null;
	    }

	    // If there isn't a thread on the queue, just return null.
	    if (threads.isEmpty()) 
		return null;

	    // The next thread on the queue should become my
	    // owner. Pop it, let it acquire me, and return it.
	    owner = pop();
	    owner.acquire(this);
	    return owner.thread;
	}

	/**
	 * Return the next thread that <tt>nextThread()</tt> would return,
	 * without modifying the state of this queue.
	 *
	 * @return	the next thread that <tt>nextThread()</tt> would
	 *		return.
	 */
	protected ThreadState pickNextThread() {
	    return isEmpty() ? null : threads.get(0);
	}
	
	public void print() {

	    Lib.assertTrue(Machine.interrupt().disabled());
	    
	    if (owner != null) {
		System.out.printf(" %5s %30s: %d/%d (%d)%n",
				  "owner", owner.thread, owner.getEffectivePriority(),
				  owner.entryTime,
				  owner.getPriority());
	    }
	    for (int i = 0; i < threads.size(); i++) {
		ThreadState ts = threads.get(i);
		System.out.printf("  %4d %30s: %d/%d (%d)%n",
				  i, ts.thread,
				  ts.getEffectivePriority(),
				  ts.entryTime,
				  ts.getPriority());
	    }

	}

	private int parent(int i) {
	    return i / 2;
	}

	private int left(int i) {
	    return 2 * i;
	}

	private int right(int i) {
	    return 2 * i + 1;
	}

	private void heapify(int i) {
	    int l = left(i);
	    int r = right(i);

	    // Set largest to be the index of the largest key among
	    // myself, my left child, and my right child.
	    int largest = 
		(l <= last() && greaterThan(l, i) ? l : 
		 r <= last() && greaterThan(r, l) ? r :
		 i);

	    // If one of my children is larger than me, swap it with
	    // me and heapify it.
	    if (largest != i) {
		swap(i, largest);
		heapify(largest);
	    }		    
	}
	
	private int last() {
	    return threads.size() - 1;
	}

	/**
	 * Every time we add a new thread to the queue, we increment
	 * this timer and associate the current value of time with the
	 * newly added thread. This allows us to remember when each
	 * thread was added, so if a thread moves up or down with
	 * priority we can know when it entered compared to other
	 * threads of the same priority.
	 */
	private int time = 0;

	/**
	 * <tt>true</tt> if this queue should transfer priority from waiting
	 * threads to the owning thread.
	 */
	public boolean transferPriority;

	/**
	 * We store the threads currently in the queue in this
	 * list. It's arranged as a binary heap.
	 */
	private ArrayList<ThreadState> threads = new ArrayList<ThreadState>();

	/**
	 * The last thread that was popped off this queue, and is
	 * currently blocking other threads on the queue.
	 */
	private ThreadState owner = null;

	/**
	 * Return true if there are no threads on the queue, false
	 * otherwise.
	 */
	public boolean isEmpty() {
	    return threads.isEmpty();
	}


	/**
	 * Returns the highest effective priority of any thread on the
	 * queue, or priorityMinimum if there are no threads on the
	 * queue.
	 */
	public int highestPriority() {
	    return isEmpty() ? priorityMinimum :
		pickNextThread().getEffectivePriority();
	}

    }

    /**
     * The scheduling state of a thread. This should include the thread's
     * priority, its effective priority, any objects it owns, and the queue
     * it's waiting for, if any.
     *
     * @see	nachos.threads.KThread#schedulingState
     */
    protected class ThreadState implements Comparable<ThreadState> {
	/**
	 * Allocate a new <tt>ThreadState</tt> object and associate it with the
	 * specified thread.
	 *
	 * @param	thread	the thread this state belongs to.
	 */
	public ThreadState(KThread thread) {
	    this.thread = thread;
	    
	    setPriority(priorityDefault);
	}

	/**
	 * Return a negative number if the other thread should be
	 * released before me, or a positive number if I should be
	 * released before the other thread. This should never return
	 * 0.
	 */
	public int compareTo(ThreadState other) {
	    // If one of the threads has a higher priority, base the
	    // decision on that. Otherwise base it on entry time.
	    int priorityDiff = getEffectivePriority() -
		other.getEffectivePriority();
	    return priorityDiff != 0 ? priorityDiff :
		other.entryTime - this.entryTime;
	}

	/**
	 * Return the priority of the associated thread.
	 *
	 * @return	the priority of the associated thread.
	 */
	public int getPriority() {
	    return priority;
	}

	/**
	 * Return the effective priority of the associated thread.
	 *
	 * @return	the effective priority of the associated thread.
	 */
	public int getEffectivePriority() {
	    return effectivePriority;
	}

	/**
	 * Set the priority of the associated thread to the specified value.
	 *
	 * @param	priority	the new priority.
	 */
	public void setPriority(int priority) {
	    // If it's already set to this priority, we don't need to do anything.
	    if (this.priority == priority)
		return;

	    this.priority = priority;

	    // If this priority is higher than my effectivePriority,
	    // we need to adjust effectivePriority.
	    if (priority > effectivePriority)
		effectivePriority = priority;

	    // If I'm on a queue at the moment, I need to notify the
	    // queue that my priority has changed.
	    if (waitQueue != null) {
		if (priority > this.priority)
		    waitQueue.priorityIncreased(this);
		else
		    waitQueue.priorityDecreased(this);
	    }
	}

	/**
	 * Called when <tt>waitForAccess(thread)</tt> (where <tt>thread</tt> is
	 * the associated thread) is invoked on the specified priority queue.
	 * The associated thread is therefore waiting for access to the
	 * resource guarded by <tt>waitQueue</tt>. This method is only called
	 * if the associated thread cannot immediately obtain access.
	 *
	 * @param	waitQueue	the queue that the associated thread is
	 *				now waiting on.
	 *
	 * @see	nachos.threads.ThreadQueue#waitForAccess
	 */
	public void waitForAccess(PriorityQueue waitQueue, int index, int entryTime) {
	    this.waitQueue = waitQueue;
	    this.index = index;
	    this.entryTime = entryTime;
	    if (ownedQueues.remove(waitQueue)) {
		calculateEffectivePriority();
	    }
	}

	/**
	 * Called when the associated thread has acquired access to whatever is
	 * guarded by <tt>waitQueue</tt>. This can occur either as a result of
	 * <tt>acquire(thread)</tt> being invoked on <tt>waitQueue</tt> (where
	 * <tt>thread</tt> is the associated thread), or as a result of
	 * <tt>nextThread()</tt> being invoked on <tt>waitQueue</tt>.
	 *
	 * @see	nachos.threads.ThreadQueue#acquire
	 * @see	nachos.threads.ThreadQueue#nextThread
	 */
	public void acquire(PriorityQueue waitQueue) {

	    Lib.assertTrue(Machine.interrupt().disabled());
	    ownedQueues.add(waitQueue);

	    this.waitQueue = null;
	}	

	void calculateEffectivePriority() {
	    effectivePriority = priority;
	    for (PriorityQueue q : ownedQueues) {
		if (q.transferPriority &&
		    q.highestPriority() > effectivePriority) {
		    effectivePriority = q.highestPriority();
		}
	    }
	}

	/** 
	 * The queue that this thread is waiting on, or null if it's
	 * active.
	 */
	public PriorityQueue waitQueue;

	/**
	 * The time that this thread entered waitQueue, or undefined
	 * if waitQueue is null.
	 */
	public int entryTime;

	/**
	 * This thread's index in waitQueue, or undefined if waitQueue
	 * is null. We need to store the index here so that
	 * increasePriority and decreasePriority can know where in the
	 * list this ThreadState object is.
	 */
	public int index;

	/**
	 * Set of queues that this thread currently owns. We use this
	 * to determine whether there is a higher priority thread that
	 * is being blocked by this thread.
	 */
	public Set<PriorityQueue> ownedQueues = new HashSet<PriorityQueue>();

	/** The thread with which this object is associated. */	   
	protected KThread thread;
	/** The priority of the associated thread. */
	protected int priority;

	protected int effectivePriority = priorityDefault;
    }

    public static void selfTest() {
	selfTestPriorityDonation();
	selfTestDecreasePriority();
    }

    /**
     * I have been using this test to make sure that PriorityScheduler
     * handles the priority inversion problem. I kick off a low
     * priority thread that acquires a lock, and then a medium
     * priority thread, and then a high priority thread which tries to
     * acquire the same lock that's held by the low priority
     * thread. If we're handling this correctly, the low thread should
     * run before the medium thread so that it frees up the resource
     * needed by the high priority thread.
     */
    public static void selfTestPriorityDonation() {
	System.out.println("Testing priority donation");
	final Lock lock = new Lock();
	
	// The low, medium, and high priority threads will append
	// characters to this StringBuilder so we can examine the
	// order of execution.
	final StringBuilder output = new StringBuilder();

	final KThread h = new KThread(new Runnable() {
		public void run() {
		    lock.acquire();
		    for (int i = 0; i < 5; i++) {
			output.append("h");
			KThread.currentThread().yield();
		    }
		    lock.release();
		}
	    });
	
	final KThread m = new KThread(new Runnable() {
		public void run() {
		    for (int i = 0; i < 5; i++) {
			output.append("m");
			if (i == 2) {
			    h.fork();
			}
			KThread.currentThread().yield();
		    }
		}
	    });
	
	final KThread l = new KThread(new Runnable() {
		public void run() {
		    lock.acquire();
		    for (int i = 0; i < 5; i++) {
			output.append("l");
			if (i == 2) {
			    m.fork();
			}
			KThread.currentThread().yield();
		    }
		    lock.release();
		}
	    });
	l.setName("L");
	m.setName("M");
	h.setName("H");

	boolean intStatus = Machine.interrupt().disable();
	ThreadedKernel.scheduler.setPriority(l, 1);
	ThreadedKernel.scheduler.setPriority(m, 2);
	ThreadedKernel.scheduler.setPriority(h, 3);
	Machine.interrupt().restore(intStatus);

	l.fork();

	// Join with the threads we kicked off to wait for them to
	// finish.
	h.join();
	l.join();
	m.join();

	// Low should loop three times and then fork middle, which
	// should preempt low and loop three times. Then medium should
	// fork high, which should donate its priority to low. Low
	// should finish looping two more times, then high should loop
	// five times, then medium should finish up.
	String expectedOut = "lllmmmllhhhhhmm";
	Lib.assertTrue(expectedOut.equals(output.toString()));
	System.out.println("High donated its priority to low: " + output);
    }

    public static void selfTestDecreasePriority() {
	System.out.println("Testing priority donation");
	final Lock lock = new Lock();
	
	// The low, medium, and high priority threads will append
	// characters to this StringBuilder so we can examine the
	// order of execution.
	final StringBuilder output = new StringBuilder();

	final KThread h = new KThread(new Runnable() {
		public void run() {
		    lock.acquire();
		    for (int i = 0; i < 5; i++) {
			output.append("h");
			KThread.currentThread().yield();
		    }
		    lock.release();
		}
	    });
	
	final KThread m = new KThread(new Runnable() {
		public void run() {
		    for (int i = 0; i < 5; i++) {
			output.append("m");
			if (i == 2) {
			    h.fork();
			}
			KThread.currentThread().yield();
		    }
		}
	    });
	
	final KThread l = new KThread(new Runnable() {
		public void run() {
		    lock.acquire();
		    for (int i = 0; i < 5; i++) {
			output.append("l");
			if (i == 1) {
			    m.fork();
			}
			if (i == 2) {
			    boolean intStatus = Machine.interrupt().disable();
			    ThreadedKernel.scheduler.setPriority(h, 1);
			    Machine.interrupt().restore(intStatus);
			}
			KThread.currentThread().yield();
		    }
		    lock.release();
		}
	    });
	l.setName("L");
	m.setName("M");
	h.setName("H");

	boolean intStatus = Machine.interrupt().disable();
	ThreadedKernel.scheduler.setPriority(l, 1);
	ThreadedKernel.scheduler.setPriority(m, 2);
	ThreadedKernel.scheduler.setPriority(h, 3);
	Machine.interrupt().restore(intStatus);

	l.fork();

	// Join with the threads we kicked off to wait for them to
	// finish.
	h.join();
	l.join();
	m.join();

	// Low should loop two times and then fork middle, which
	// should preempt low and loop three times. Then medium should
	// fork high, which should donate its priority to low. Low
	// should loop once, and then lower the priority of high. This
	// should cause medium to once again be the highest priority
	// thread, so it will loop two more times. THen low loops two
	// more times, releases the lock, and high loops five times.
	// five times, then medium should finish up.
	String expectedOut = "llmmmlmmllhhhhh";
	Lib.assertTrue(expectedOut.equals(output.toString()),
		       "Got output " + output);
	System.out.println("High donated its priority to low: " + output);
    }

}
